# Copyright 2021-2022 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#

# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Read information from a xyz format file.
"""

import numpy as np
from collections import namedtuple


def xyzparser(file_name: str):
    """Read a xyz file and return atom information with numpy array format.
    Args:
        file_name(str): The xyz file name, absolute path is suggested.
    Returns:
        atom_names(np.str_): Numpy array contain all atom names.
        crds(np.float32): The numpy array format of coordinates.
    """
    xyz_obj = namedtuple('XYZObject', ['atom_names', 'crds'])
    with open(file_name) as f:
        lines = f.readlines()

    total_num = int(lines[0].strip())
    atoms = []
    crd = []
    for line in lines[2:2+total_num]:
        data = line.split()
        atoms.append(data[0])
        crd.append(data[1:])

    atom_names = np.array(atoms, np.str_).flatten()
    crds = np.array(crd, np.float32)

    return xyz_obj(atom_names, crds)